#!/usr/bin/env python
# -*- coding: utf-8 -*-
'''
# @File   : trt_test.py
# @Author : yuanwenjin
# @Mail   : xxxx@mail.com
# @Date   : 2020/07/01 12:28:03
# @Docs   : 对tensorRT进行测试, tensorflow转tensorRT
'''

import os
import numpy as np
from trt_infer import trt_infer
from PIL import Image
from PIL import ImageFile
ImageFile.LOAD_TRUNCATED_IMAGES = True

def image_preprocess(img, mask256, reshaped_size=None, rgb_means=(123.68, 116.78, 103.94)):
    """图像预处理
    ### Args:
        - img: H*W*C, PIL image, rgb
        - mask: H*W*C, array, same with reshaped
        - reshaped_size, 2*1, (height, width), tuple
        - rgb_means: C*1, (meanR, meanG, meanB), tuple
    ### Returns: CHW
    """
    img = np.array(img.resize(reshaped_size, Image.ANTIALIAS), dtype='float32')
    img = img * mask256
    img[:, :, 0] = (img[:, :, 0] - rgb_means[0])
    img[:, :, 1] = (img[:, :, 1] - rgb_means[1])
    img[:, :, 2] = (img[:, :, 2] - rgb_means[2])

    return img.transpose([2, 0, 1])

if __name__ == "__main__":

    mask = Image.open('mask/mask.png')
    mask256 = mask.resize((256, 256))
    mask256 = np.array(mask256, 'uint8') > 0
    x = np.expand_dims(mask256, axis=2)
    mask256 = np.concatenate([x, x, x], axis=2)  

    image_dir = './test_images_3/'
    imgs_file = [os.path.join(image_dir, f) for f in os.listdir(image_dir)]

    img_num = len(imgs_file)
    # img_num = 35
    images = [Image.open(img_bin).resize((480, 480)) for img_bin in imgs_file]

    engine = trt_infer('model/resnet_v2_152_s256_c15_frozen_model_b16_fp32.engine', batch_size=16)
    step = engine.batch_size
    class_num = 15

    preds = np.empty([0, class_num])
    for i in range(0, img_num, step):    # 将图像按照BatchSize分批处理
        imgs = images[i:min(i+step, img_num)]
        imgs = [image_preprocess(img, mask256, (256, 256)) for img in imgs]
        prediction = engine.infer(imgs)
        prediction_reshape = prediction[0].reshape(-1, class_num)
        preds = np.concatenate((preds, prediction_reshape), axis=0)

    gailv = np.array(preds)
    predictions = np.argmax(gailv, 1).tolist()
    for idx, pre in enumerate(predictions):
        if gailv[idx][pre] < 0.8 and (pre < 11 and pre > 2):
            predictions[idx] = 11

    for idx, img in enumerate(imgs_file[:img_num]):
        print('{} {}'.format(img, predictions[idx]))
